package com.zyh.day05.function

import org.apache.spark.sql.expressions.{MutableAggregationBuffer, UserDefinedAggregateFunction}
import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, LongType, StructType}
import org.apache.spark.sql.{DataFrame, Row, SparkSession}

object UserDefinedAggregateFunctionTest {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession.builder()
      .master("local[*]")
      .appName("udaf")
      .getOrCreate()

    import spark.implicits._
    val list = List((1, "xiao1hei", 18), (2, "xiao2hei", 20), (3, "xiao3hei", 22))
    val df: DataFrame = list.toDF("id","name","age")
    //定义一个函数 avg_age 求平均值
    spark.udf.register("avg_age",new AvgAgeFunction)
    df.createOrReplaceTempView("t_user")
    val sql =
      """
        |select avg_age(age)
        |from t_user
        |""".stripMargin
    val result: DataFrame = spark.sql(sql)
    result.show()
    spark.close()
  }
}
//自定义组函数
class AvgAgeFunction extends UserDefinedAggregateFunction{
  //返回 自定义函数接受的参数的类型信息
  override def inputSchema: StructType = new StructType().add("age",IntegerType)

  //返回聚合函数中间结果的类型信息 总和与总个数
  override def bufferSchema: StructType = new StructType().add("sum",DoubleType).add("count",LongType)

  //返回 自定义函数的返回值类型信息
  override def dataType: DataType = DoubleType

  //用于设定面对相同输入参数时 自定义函数是否返回相同的结果
  override def deterministic: Boolean = true

  //用来初始化中间结果
  override def initialize(buffer: MutableAggregationBuffer): Unit = {
    buffer(0) = 0.0
    buffer(1) = 0L
  }

  //input表示自定义函数接受的一行参数,在这里表示接受的一个年龄
  //buffer表示中间结果
  override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
    buffer(0) = buffer.getDouble(0)+input.getInt(0)
    buffer(1) = buffer.getLong(1)+1
  }

  //合并方法 将各个分区结果进行合并
  override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
    buffer1(0) = buffer1.getDouble(0) + buffer2.getDouble(0)
    buffer1(1) = buffer1.getLong(1) + buffer2.getLong(1)
  }

  //返回最终的结果
  override def evaluate(buffer: Row): Any = {
    buffer.getDouble(0)/buffer.getLong(1)
  }
}